from flask import Blueprint, render_template, redirect, url_for, flash, request, session
from app import db
from app.models.user import User
from app.models.role import Role
from app.models.menu import Menu

auth_bp = Blueprint('auth', __name__)

@auth_bp.route('/')
def main():
    # 如果用户已登录，渲染base.html
    if 'user_id' in session:
        # 每次请求都从数据库获取最新菜单数据
        role_id = session['role_id']
        role = Role.query.get(role_id)
        menus = role.menus
        
        # 构建菜单树结构
        menu_dict = {}
        menu_list = []
        
        # 首先将所有菜单添加到字典中
        for menu in menus:
            menu_dict[menu.id] = {
                'id': menu.id,
                'title': menu.name,
                'url': menu.url,
                'icon': menu.icon,
                'parent_id': menu.parent_id,
                'children': []
            }
        
        # 构建父子关系
        for menu_id, menu in menu_dict.items():
            if menu['parent_id'] is None:
                # 没有父菜单的是顶级菜单
                menu_list.append(menu)
            elif menu['parent_id'] in menu_dict:
                # 将子菜单添加到父菜单中
                menu_dict[menu['parent_id']]['children'].append(menu)
        
        # 更新session中的菜单数据
        session['menus'] = menu_list
        
        # 传递菜单数据给模板
        return render_template('layout/main.html', menuData=menu_list, username=session['username'])
    # 否则重定向到登录页面
    return redirect(url_for('auth.login'))

@auth_bp.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']

        user = User.query.filter_by(username=username, password=password).first()
        if user:
            # 登录成功，存储用户信息到session
            session['user_id'] = user.id
            session['username'] = user.username
            session['role_id'] = user.role_id

            # 获取用户角色
            role = Role.query.get(user.role_id)
            # 获取角色对应的菜单
            menus = role.menus
            
            # 构建菜单树结构
            menu_dict = {}
            menu_list = []
            
            # 首先将所有菜单添加到字典中
            for menu in menus:
                menu_dict[menu.id] = {
                    'id': menu.id,
                    'title': menu.name,
                    'url': menu.url,
                    'icon': menu.icon,
                    'parent_id': menu.parent_id,
                    'children': []
                }
            
            # 构建父子关系
            for menu_id, menu in menu_dict.items():
                if menu['parent_id'] is None:
                    # 没有父菜单的是顶级菜单
                    menu_list.append(menu)
                elif menu['parent_id'] in menu_dict:
                    # 将子菜单添加到父菜单中
                    menu_dict[menu['parent_id']]['children'].append(menu)
            
            # 存储构建好的菜单树到session
            session['menus'] = menu_list

            flash('登录成功！', 'success')
            return redirect(url_for('auth.main'))
        else:
            flash('用户名或密码错误！', 'error')
            return redirect(url_for('auth.login'))
    return render_template('auth/login.html')

@auth_bp.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form['username']
        password = request.form['password']
        confirm_password = request.form['confirm_password']

        if password != confirm_password:
            flash('两次密码输入不一致！', 'error')
            return redirect(url_for('auth.register'))

        existing_user = User.query.filter_by(username=username).first()
        if existing_user:
            flash('用户名已存在！', 'error')
            return redirect(url_for('auth.register'))

        # 获取默认角色（假设id为2的角色是普通用户）
        default_role = Role.query.get(2)
        if not default_role:
            default_role = Role(name='user', description='普通用户')
            db.session.add(default_role)
            db.session.commit()

        new_user = User(username=username, password=password, role_id=default_role.id)
        db.session.add(new_user)
        db.session.commit()

        flash('注册成功，请登录！', 'success')
        return redirect(url_for('auth.login'))
    return render_template('auth/register.html')

@auth_bp.route('/logout')
def logout():
    # 清空session
    session.clear()
    flash('已退出登录！', 'success')
    return redirect(url_for('auth.login'))

# 登录检查装饰器
def login_required(f):
    from functools import wraps

    @wraps(f)
    def decorated_function(*args, **kwargs):
        if 'user_id' not in session:
            flash('请先登录！', 'error')
            return redirect(url_for('auth.login'))
        return f(*args, **kwargs)

    return decorated_function